#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <string.h>
#include <mpi.h>
#include "blowfish.h"

extern uint32_t P[P_SIZE];

extern uint32_t S1[S_SIZE];
extern uint32_t S2[S_SIZE];
extern uint32_t S3[S_SIZE];
extern uint32_t S4[S_SIZE];

int main(int argc, char* argv[])
{
	FILE *keyfile, *infile, *outfile;
	uint32_t *key;
	uint64_t *file;
	struct stat buf;
	int i;
	int file_size, key_size, blocks, padded_size;
	char* key_str,*outfile_name;
	char rem_bytes;
	uint64_t *working_set;
	int working_set_size, rank, num_threads;
	int *working_set_sizes, *offsets;
	//MPI_Status status;
	
	padded_size = 0;
	file = NULL;
	key_str = NULL;
	blocks = 0;
	outfile = NULL;
	infile = NULL;
	keyfile = NULL;
	
	MPI_Init(&argc,&argv);
	MPI_Comm_rank(MPI_COMM_WORLD,&rank);
	MPI_Comm_size(MPI_COMM_WORLD,&num_threads);
	
	working_set_sizes = (int*)calloc(num_threads,sizeof(int));
	offsets = (int*)calloc(num_threads,sizeof(int));
	
	if(rank == 0)
	{
		if(argc != 3)
		{
			printf("Usage: ./bfencrypt <filename> <keyfile>\n");
			exit(1);
		}
		
		infile = fopen(argv[1], "rb");
		
		if(infile == NULL)
		{
			printf("Could not open input file %s.\n",argv[1]);
			exit(1);
		}
		
		keyfile = fopen(argv[2], "r");
		
		if(keyfile == NULL)
		{
			printf("Could not open keyfile.\n");
			exit(1);
		}
		
		outfile_name = (char*)malloc(strlen(argv[1])+5);
		strcpy(outfile_name,argv[1]);
		strcat(outfile_name,".blf");
		
		outfile = fopen(outfile_name,"wb");
		free(outfile_name);
		
		if(outfile == NULL)
		{
			printf("Could not open output file %s.\n",outfile_name);
			exit(1);
		}
		
		fstat(fileno(keyfile), &buf);
		
		file_size = buf.st_size;
		
		blocks = file_size/8;
		
		key_size = blocks*8*sizeof(uint32_t);
		
		//printf("Key file size: %d\nKey size: %d\nBlocks: %d\n",file_size,key_size,blocks);
		
		key_str = (char*)calloc(17,sizeof(char));
	}
	
	MPI_Bcast(&key_size,1,MPI_INT,0,MPI_COMM_WORLD);
	key = (uint32_t*)malloc(key_size/8);
	
	if(rank == 0)
	{
		
		for(i = 0; i < 8*blocks; i+=8)
		{
			fread((void*)(key_str),1,8,keyfile);
			//printf("%d: %s\n",i/8,key_str);
			key[i/8] = (uint32_t)strtoul(key_str,NULL,16);
		}
		
		/*for(i = 0; i < blocks; i++)
		{
			printf("%#lx\n",key[i]);
		}*/
		
		gen_subkeys(key,key_size);
		
		fstat(fileno(infile), &buf);
		
		file_size = buf.st_size;
		
		if(file_size % 8)
		{
			padded_size = ((file_size/sizeof(uint64_t))+1)*sizeof(uint64_t);
		}
		else
		{
			padded_size = file_size;
		}
		
		working_set_size = file_size/sizeof(uint64_t)/num_threads;
		
		working_set_sizes[0] = ((file_size % (sizeof(uint64_t)*num_threads)) > 0) + ((file_size % (sizeof(uint64_t)*num_threads)) > 8);
		
		for(i = 0; i < num_threads; i++)
		{
			working_set_sizes[i] += working_set_size;
			offsets[i+1] = offsets[i] + working_set_sizes[i];
		}
		
		file = (uint64_t*)calloc(padded_size/sizeof(uint64_t),sizeof(uint64_t));
		//file = memset(file,0,padded_size);
		
		//printf("Infile size: %d\n",file_size);
		
		if(file_size % 8)
		{
			rem_bytes = ((file_size/8)+1)*8 - file_size;
		}
		else
		{
			rem_bytes = 0;
		}
		//printf("%d bytes padded.\n",rem_bytes);
		/*for(i = 0; i < file_size; i+=8)
		{
			chunk = 0;
			//fread((void*)(&chunk),1,8,infile);
			//blowfish(&chunk);
			//fwrite((void*)(&chunk),1,8,outfile);
		}*/
		fread((void*)file,1,file_size,infile);
		
	}
	//MPI_Barrier(MPI_COMM_WORLD);
	//MPI_Bcast(&working_set_size,1,MPI_UNSIGNED_LONG_LONG,0,MPI_COMM_WORLD);
	//MPI_Bcast(&key_size,1,MPI_INT,0,MPI_COMM_WORLD);
	//MPI_Bcast(key,key_size/32,MPI_UNSIGNED,0,MPI_COMM_WORLD);
	MPI_Bcast(P,P_SIZE,MPI_UNSIGNED,0,MPI_COMM_WORLD);
	MPI_Bcast(S1,S_SIZE,MPI_UNSIGNED,0,MPI_COMM_WORLD);
	MPI_Bcast(S2,S_SIZE,MPI_UNSIGNED,0,MPI_COMM_WORLD);
	MPI_Bcast(S3,S_SIZE,MPI_UNSIGNED,0,MPI_COMM_WORLD);
	MPI_Bcast(S4,S_SIZE,MPI_UNSIGNED,0,MPI_COMM_WORLD);
	
	MPI_Scatter(working_set_sizes,1,MPI_INT,&working_set_size,1,MPI_INT,0,MPI_COMM_WORLD);
	working_set = (uint64_t*)malloc(working_set_size*sizeof(uint64_t));
	//MPI_Scatter(file,working_set_size,MPI_UNSIGNED_LONG_LONG,working_set,working_set_size,MPI_UNSIGNED_LONG_LONG,0,MPI_COMM_WORLD);
	MPI_Scatterv(file,working_set_sizes,offsets,MPI_UNSIGNED_LONG_LONG,working_set,working_set_size,MPI_UNSIGNED_LONG_LONG,0,MPI_COMM_WORLD);
	
	for(i = 0; i < working_set_size; i++)
	{
		blowfish(working_set+i);
	}
	
	MPI_Gatherv(working_set,working_set_size,MPI_UNSIGNED_LONG_LONG,file,working_set_sizes,offsets,MPI_UNSIGNED_LONG_LONG,0,MPI_COMM_WORLD);
	
	if(rank == 0)
	{
		fwrite((void*)file,1,padded_size,outfile);
		fwrite((void*)(&rem_bytes),1,1,outfile);
		
		fclose(keyfile);
		fclose(infile);
		fclose(outfile);
	}
	
	MPI_Finalize();
	return 0;
}
